import sys
import cv2 
from merge_pic import mergeImge,getMachine
from PyQt5.QtWidgets import  QApplication,QPushButton,QFileDialog,QMainWindow,QMessageBox,QLabel
from PyQt5.QtGui import QIcon,QImage, QPixmap

class myGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.img1 = None
        self.img2 = None
        self.outimg = None
        self.outimg_state = 0
        self.imgNum = 0
        self.sift,self.flann = getMachine()
        self.initUI()

    def initUI(self): 
        self.setFixedSize(1000,800)                    
        self.setWindowTitle('IMAGE MERGE')  
        #self.statusBar()
        self.setWindowIcon(QIcon('./source/imgsrc/icon.png')) 
        self.imglbl = QLabel(self)
        #self.imglbl.setScaledContents (True)  
        self.imglbl.resize(900,700)
        self.imglbl.move(50,50)

        choiceImg = QPushButton('加载图片', self)  
        choiceImg.setFixedSize(100,50)
        choiceImg.move((1000//3-100)//2, 30) 
        choiceImg.clicked.connect(self.openFile)


        mergeImg = QPushButton('拼接图片', self)  
        mergeImg.setFixedSize(100,50)
        mergeImg.move((1000//3-100)//2+1000//3, 30) 
        mergeImg.clicked.connect(self.merge)

        saveImg = QPushButton('保存图片', self)  
        saveImg.setFixedSize(100,50)
        saveImg.move((1000//3-100)//2+1000//3*2, 30) 
        saveImg.clicked.connect(self.saveFile)

        self.statusBar().showMessage("请加载图片!!!")
        

        self.show() 
    def saveFile(self):
        if self.outimg_state:
            filter = "Images (*.jpg);;Images (*.bmp);;Images (*.png)"
            fname = QFileDialog.getSaveFileName(self, 'Save file', './output/',filter)
            cv2.imwrite(fname[0],self.outimg)
        else:
            self.statusBar().showMessage("没有拼接成功的图片!!!")
    def merge(self):
        if self.imgNum == 2:
            self.statusBar().showMessage("正在拼接中，请耐心等待~~~")
            QApplication.processEvents()
            res = mergeImge(self.img1,self.img2,self.sift,self.flann)
            self.img2 = None
            if not res[0]:
                self.statusBar().showMessage("没有足够的特征点,拼接失败!!!")
                self.imgNum = 0
                self.img1 = None
            else:
                self.imgNum = 1
                self.img1 = res[1]
                self.outimg = res[2]
                outimg = self.outimg.copy()
                self.outimg_state=1
                if outimg.shape[1] >900:
                    outimg = cv2.resize(outimg,(900,int(900/outimg.shape[1] * outimg.shape[0])))
                if outimg.shape[0] >700:
                    outimg = cv2.resize(outimg,(int(700/outimg.shape[0] * outimg.shape[1]),700))
                imgrgb = cv2.cvtColor(outimg,cv2.COLOR_BGR2RGB)
                w = imgrgb.shape[1]  # 获取图像大小
                h = imgrgb.shape[0]
                frame = QImage(imgrgb.data, w,h,w*3,QImage.Format_RGB888)
                pix = QPixmap.fromImage(frame)
                self.imglbl.setPixmap (pix)
                self.statusBar().showMessage("继续拼接-->请加载新的图片")
        else:
            self.statusBar().showMessage("图片数量不足，请继续添加...")

    def openFile(self):
        fname = QFileDialog.getOpenFileName(self, 'Choice Image', './img/')
        if not self.imgNum:
            self.img1 = cv2.imread(fname[0])
            if self.img1 is None:
                self.statusBar().showMessage("加载图片失败,请重新选择!!!")
                QMessageBox.warning(self,'Warning', '无效的图片!   ')
            else:
                self.imgNum += 1
                self.statusBar().showMessage("还需加载一张图片")
        elif self.imgNum==1:
            self.img2 = cv2.imread(fname[0])
            if self.img2 is None:
                self.statusBar().showMessage("加载图片失败,请重新选择!!!")
                QMessageBox.warning(self,'Warning', '无效的图片!   ')
            else:
                self.imgNum += 1
                self.statusBar().showMessage("已加载两张图片，可以拼接!!!")   
if __name__ == '__main__': 
    app = QApplication(sys.argv)

    ex = myGUI()
    sys.exit(app.exec_())  